package com.example.ai.agent;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.example.ai.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具具体实现了think act 的方法
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    //可用工具列表
    private  final ToolCallback[] availableTools;
    //保存工具聊天响应
    private ChatResponse toolCallChatResponse;
    //工具管理器
    private final ToolCallingManager toolCallManager;

    //禁用以及自己实现工具调用
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallManager = ToolCallingManager.builder().build();
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();

    }



    @Override
    public boolean think() {
        if (getNextStepPrompt() != null && !getNextStepPrompt().isEmpty()){
            UserMessage userMessage = new UserMessage(getNextStepPrompt());
            getMessagesList().add(userMessage);
        }
        List<Message> messageList = getMessagesList();
        Prompt prompt =new Prompt(messageList, chatOptions);
        try {

            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();


        this.toolCallChatResponse = chatResponse;
        AssistantMessage assistantMessage = null;
        if (chatResponse != null){
            assistantMessage = chatResponse.getResult().getOutput();
        }
        String result = null;
        if (assistantMessage != null){
            result = assistantMessage.getText();
        }
        log.info("{}的思考: {}", getName(), result);
        List<AssistantMessage.ToolCall> toolCallList = null;
        if (assistantMessage != null){
            toolCallList = assistantMessage.getToolCalls();
            log.info("{} 选择了 {} 个工具调用",getName() ,toolCallList.size());
        }
        String toolCallInfo = toolCallList.stream()
                .map(toolCall -> String.format("工具名称: %s, 参数 %s ",toolCall.name(),toolCall.arguments()))
                .collect(Collectors.joining("\n"));
        log.info("工具调用信息: {}", toolCallInfo);
        if (toolCallList.isEmpty()){
            getMessagesList().add(assistantMessage);
            return false;
        }else {
            return true;
        }
    }catch (Exception e){
        log.error("{}的思考失败{}",getName(),e.getMessage());
        getMessagesList().add(new AssistantMessage("思考失败"+e.getMessage()));
        return false;
    }
    }

    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()){
            return "工具词用";
        }
        Prompt prompt = new Prompt(getMessagesList(), chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallManager.executeToolCalls(prompt, toolCallChatResponse);
        //记录上下文
        setMessagesList(toolExecutionResult.conversationHistory());
        //返回当前结果
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());
        String result = toolResponseMessage.getResponses().stream()
                .map(response -> "工具" +response.name() + "完成任务 结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        if (toolResponseMessage.getResponses().stream().anyMatch(response -> response.name().equals("doTerminate"))){
            setState(AgentState.FINISHED);
            return "工具调用失败";
        }
        log.info("工具词调用结果：\n{}", result);
        return result;
    }
    @Override
    public String step() {
        return super.step();
    }

}
